package cn.lovemmd.model.util;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Refill;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 速率限制工具类
 * 用于防止API被滥用
 * 
 * @author TLS
 * @since 2024-01-01
 */
@Component
public class RateLimitUtil {

    private final Map<String, Bucket> buckets = new ConcurrentHashMap<>();

    /**
     * 获取IP对应的速率限制桶
     * 
     * @param ip IP地址
     * @return 速率限制桶
     */
    public Bucket getBucket(String ip) {
        return buckets.computeIfAbsent(ip, this::createBucket);
    }

    /**
     * 创建速率限制桶
     * 
     * @param ip IP地址
     * @return 速率限制桶
     */
    private Bucket createBucket(String ip) {
        // 每分钟最多10次请求
        Bandwidth limit = Bandwidth.classic(10, Refill.greedy(10, Duration.ofMinutes(1)));
        return Bucket.builder()
                .addLimit(limit)
                .build();
    }

    /**
     * 检查是否允许请求
     * 
     * @param ip IP地址
     * @return 是否允许
     */
    public boolean tryConsume(String ip) {
        Bucket bucket = getBucket(ip);
        return bucket.tryConsume(1);
    }

    /**
     * 获取剩余令牌数
     * 
     * @param ip IP地址
     * @return 剩余令牌数
     */
    public long getAvailableTokens(String ip) {
        Bucket bucket = getBucket(ip);
        return bucket.getAvailableTokens();
    }

    /**
     * 清理过期的桶
     */
    public void cleanup() {
        // 这里可以实现清理逻辑，比如定期清理长时间未使用的桶
        // 为了简化，暂时不实现
    }
}
